import sqlite3
import psycopg2
from psycopg2 import sql
from pymongo import MongoClient
import bson
import hashlib

# 如果单独运行此python程序，需要提前在postgresql中建立名为'bookstore'的数据库

def migrate_data():
    try:
        # SQLite 连接
        sqlite_conn = sqlite3.connect("fe/data/book.db")
        sqlite_conn.text_factory = bytes  # 强制所有文本字段（包括 picture）以原始字节方式返回
        sqlite_cursor = sqlite_conn.cursor()

        # PostgreSQL 连接
        postgres_conn_info = {
            'dbname': 'bookstore',
            'user': 'postgres',
            'password': 'wyd370685.',
            'host': 'localhost',
            'port': '5432'
        }
        postgres_conn = psycopg2.connect(**postgres_conn_info)
        postgres_cursor = postgres_conn.cursor()

        # MongoDB 连接
        mongo_client = MongoClient('mongodb://localhost:27017/')
        mongo_db = mongo_client['bookstore']
        mongo_collection = mongo_db['book_details']
        mongo_collection.create_index([("book_id", 1)], unique=True)

        # 创建 PostgreSQL 表
        postgres_cursor.execute("""
            CREATE TABLE IF NOT EXISTS book (
                id TEXT PRIMARY KEY,
                title TEXT NOT NULL,
                author TEXT,
                publisher TEXT,
                original_title TEXT,
                translator TEXT,
                pub_year TEXT,
                pages INTEGER,
                price INTEGER,
                currency_unit TEXT,
                binding TEXT,
                isbn TEXT,
                tags TEXT,
                search_vector TSVECTOR
            );
        """)
        postgres_conn.commit()

        # 从 SQLite 读取数据
        sqlite_cursor.execute('SELECT * FROM book')
        rows = sqlite_cursor.fetchall()

        # 批量迁移数据
        for row in rows:
            process_row(row, mongo_collection, postgres_cursor)

        # 提交所有 PostgreSQL 事务
        postgres_conn.commit()

        print("Data Migration Completed Successfully!")

    except Exception as e:
        print(f"Error: {e}")
        if postgres_conn:
            postgres_conn.rollback()
    finally:
        # 关闭所有连接
        if sqlite_conn:
            sqlite_conn.close()
        if postgres_cursor:
            postgres_cursor.close()
        if postgres_conn:
            postgres_conn.close()
        if mongo_client:
            mongo_client.close()

def process_row(row, mongo_collection, postgres_cursor):
    (
        book_id, title, author, publisher,
        original_title, translator, pub_year,
        pages, price, currency_unit, binding,
        isbn, author_intro, book_intro,
        content, tags, picture
    ) = row

    # Debug 显示图片唯一 hash
    pic_hash = hashlib.md5(picture).hexdigest() if picture else None
    # print(f"[DEBUG] book_id={book_id}, picture_md5={pic_hash}")

    picture_data = bson.Binary(picture) if picture else None

    mongo_json = {
        "book_id": book_id,
        "author_intro": author_intro,
        "book_intro": book_intro,
        "content": content,
        "pictures": picture_data
    }

    mongo_collection.update_one(
        {"book_id": book_id},
        {"$set": mongo_json},
        upsert=True
    )

    # 处理 tags 字段，确保为字符串类型，保留所有词
    if isinstance(tags, bytes):
        tags_str = tags.decode('utf-8').strip()
    else:
        tags_str = str(tags).strip()

    postgres_cursor.execute(sql.SQL("""
        INSERT INTO book (id, title, author, publisher, original_title, 
                        translator, pub_year, pages, price, currency_unit, 
                        binding, isbn, tags)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    """), (book_id, title, author, publisher, original_title,
          translator, pub_year, pages, price, currency_unit,
          binding, isbn, tags_str))

# 执行迁移
if __name__ == "__main__":
    migrate_data()    